"""聊天图定义"""

from __future__ import annotations

from typing import Any, Dict, List, Literal, TypedDict, Optional
from enum import Enum

from langgraph.graph import StateGraph, END


class IntentType(Enum):
    """意图类型枚举"""
    QA = "qa"  # 知识问答
    CODE_ANALYSIS = "code_analysis"  # 代码分析
    ERROR_DIAGNOSIS = "error_diagnosis"  # 错误诊断
    DEEP_RESEARCH = "deep_research"  # 深度研究
    CODE_GENERATION = "code_generation"  # 代码生成
    OPTIMIZATION = "optimization"  # 代码优化
    TUTORIAL = "tutorial"  # 教程指导
    OTHER = "other"  # 其他


class ResearchStage(Enum):
    """研究阶段枚举"""
    INITIAL = "initial"  # 初始阶段
    ANALYSIS = "analysis"  # 分析阶段
    HYPOTHESIS = "hypothesis"  # 假设阶段
    VERIFICATION = "verification"  # 验证阶段
    SYNTHESIS = "synthesis"  # 综合阶段
    CONCLUSION = "conclusion"  # 结论阶段


class CodeAnalysisResult(TypedDict):
    """代码分析结果 - 增强版"""
    has_code: bool
    code_blocks: List[str]
    code_type: str  # function, struct, trait, impl, contract, module
    syntax_errors: List[Dict[str, Any]]
    ast_analysis: Optional[Dict[str, Any]]
    complexity_score: float
    dependencies: List[str]
    security_issues: List[Dict[str, Any]]
    quality_metrics: Optional[Dict[str, Any]]  # 新增：代码质量评估
    security_rating: Optional[Dict[str, Any]]  # 新增：安全评级


class ErrorDiagnosisResult(TypedDict):
    """错误诊断结果"""
    error_type: str
    error_patterns: List[str]
    root_cause: str
    fix_suggestions: List[Dict[str, Any]]
    related_docs: List[str]
    confidence_score: float


class ResearchContext(TypedDict):
    """研究上下文"""
    current_stage: ResearchStage
    research_questions: List[str]
    hypotheses: List[Dict[str, Any]]
    evidence: List[Dict[str, Any]]
    intermediate_results: List[Dict[str, Any]]
    final_report: Optional[str]


class ChatState(TypedDict):
    """增强的聊天状态定义"""
    # 基础对话信息
    messages: List[Dict[str, Any]]
    session_id: str
    user_intent: Optional[IntentType]
    confidence_score: float
    
    # Cairo 代码分析
    cairo_context: CodeAnalysisResult
    
    # 错误诊断
    error_diagnosis: Optional[ErrorDiagnosisResult]
    
    # 深度研究
    research_context: Optional[ResearchContext]
    research_depth: int
    
    # 知识检索
    retrieved_knowledge: List[Dict[str, Any]]
    knowledge_relevance_scores: List[float]
    
    # 流程控制
    next_action: Optional[str]
    loop_count: int
    max_loops: int
    
    # 输出控制
    response_format: str  # text, markdown, structured
    include_code_examples: bool
    include_references: bool


# 导入所有节点函数
from .nodes.intent_classifier import classify_intent
from .nodes.cairo_analyzer import analyze_cairo_code
from .nodes.error_analyzer import analyze_error
from .nodes.research_planner import plan_research, execute_research_stage, synthesize_research, conclude_research
from .nodes.answer_generator import generate_answer
from ..services.knowledge_service import retrieve_knowledge


def route_after_intent_classification(state: ChatState) -> str:
    """意图分类后的路由决策"""
    user_intent = state.get("user_intent")
    confidence = state.get("confidence_score", 0)
    user_text = state["messages"][-1]["content"] if state["messages"] else ""
    
    # 降低置信度阈值，让更多查询进入详细分析
    if confidence < 0.3:
        return "knowledge_retrieval"
    
    # 对于包含分析、特点、原理等关键词的查询，优先进入代码分析
    analysis_keywords = ["特点", "分析", "原理", "机制", "详细", "深入", "解释", "说明"]
    if any(keyword in user_text for keyword in analysis_keywords):
        return "code_analysis"
    
    # 根据意图路由到不同的处理分支
    if user_intent == IntentType.ERROR_DIAGNOSIS:
        return "error_analysis"
    elif user_intent == IntentType.CODE_ANALYSIS:
        return "code_analysis"
    elif user_intent == IntentType.DEEP_RESEARCH:
        return "research_planning"
    elif user_intent == IntentType.QA and confidence > 0.4:
        return "code_analysis"  # 高置信度的QA也进入代码分析
    else:
        return "code_analysis"  # 默认先分析是否有代码


def route_after_code_analysis(state: ChatState) -> str:
    """代码分析后的路由决策"""
    cairo_context = state.get("cairo_context", {})
    user_intent = state.get("user_intent")
    
    # 如果有语法错误，转到错误分析
    if cairo_context.get("syntax_errors"):
        return "error_analysis"
    
    # 如果是深度研究意图，转到研究规划
    if user_intent == IntentType.DEEP_RESEARCH:
        return "research_planning"
    
    # 否则进行知识检索
    return "knowledge_retrieval"


def route_after_error_analysis(state: ChatState) -> str:
    """错误分析后的路由决策"""
    error_diagnosis = state.get("error_diagnosis")
    user_intent = state.get("user_intent")
    
    # 如果是深度研究模式，继续深入分析
    if user_intent == IntentType.DEEP_RESEARCH and error_diagnosis:
        return "research_planning"
    
    # 否则直接生成答案
    return "answer_generation"


def route_after_knowledge_retrieval(state: ChatState) -> str:
    """知识检索后的路由决策"""
    user_intent = state.get("user_intent")
    research_depth = state.get("research_depth", 0)
    
    # 如果是深度研究模式，转到研究规划
    if user_intent == IntentType.DEEP_RESEARCH or research_depth > 1:
        return "research_planning"
    
    # 否则直接生成答案
    return "answer_generation"


def route_research_execution(state: ChatState) -> str:
    """研究执行路由决策（修复版）"""
    research_context = state.get("research_context")
    if not research_context:
        return "answer_generation"
    
    current_stage = research_context.get("current_stage", ResearchStage.INITIAL)
    loop_count = state.get("loop_count", 0)
    max_loops = state.get("max_loops", 2)  # 进一步减少最大循环次数
    
    # 检查是否达到最大循环次数
    if loop_count >= max_loops:
        return "answer_generation"  # 直接生成答案，避免过多循环
    
    # 简化阶段推进逻辑，避免循环回到knowledge_retrieval
    if current_stage == ResearchStage.INITIAL:
        # 第一次执行，进入分析阶段
        return "research_execution"  # 自循环一次进入下一阶段
    elif current_stage == ResearchStage.ANALYSIS:
        # 分析完成，直接进入综合阶段，跳过假设和验证
        return "research_synthesis"
    elif current_stage == ResearchStage.HYPOTHESIS:
        return "research_synthesis"
    elif current_stage == ResearchStage.VERIFICATION:
        return "research_synthesis"
    elif current_stage == ResearchStage.SYNTHESIS:
        return "answer_generation"
    elif current_stage == ResearchStage.CONCLUSION:
        return "answer_generation"
    else:
        return "answer_generation"


def should_continue_research(state: ChatState) -> str:
    """判断是否继续研究"""
    research_context = state.get("research_context")
    if not research_context:
        return "answer_generation"
    
    current_stage = research_context.get("current_stage", ResearchStage.INITIAL)
    loop_count = state.get("loop_count", 0)
    max_loops = state.get("max_loops", 5)
    
    # 如果已经完成或达到最大循环次数
    if current_stage == ResearchStage.CONCLUSION or loop_count >= max_loops:
        return "answer_generation"
    
    # 继续研究
    return "research_execution"


def initialize_state(state: ChatState) -> ChatState:
    """初始化状态"""
    # 设置默认值
    if "session_id" not in state:
        import uuid
        state["session_id"] = str(uuid.uuid4())
    
    if "confidence_score" not in state:
        state["confidence_score"] = 0.0
    
    if "cairo_context" not in state:
        state["cairo_context"] = {
            "has_code": False,
            "code_blocks": [],
            "code_type": "none",
            "syntax_errors": [],
            "ast_analysis": None,
            "complexity_score": 0.0,
            "dependencies": [],
            "security_issues": []
        }
    
    if "retrieved_knowledge" not in state:
        state["retrieved_knowledge"] = []
    
    if "knowledge_relevance_scores" not in state:
        state["knowledge_relevance_scores"] = []
   # 设置循环控制（优化版）
    if "loop_count" not in state:
        state["loop_count"] = 0
    
    if "max_loops" not in state:
        # 根据研究深度设置不同的循环限制
        research_depth = state.get("research_depth", 1)
        if research_depth >= 3:  # 深度研究模式
            state["max_loops"] = 3  # 减少循环次数
        else:
            state["max_loops"] = 2  # 普通模式更少循环
    
    if "response_format" not in state:
        state["response_format"] = "structured"
    
    if "include_code_examples" not in state:
        state["include_code_examples"] = True
    
    if "include_references" not in state:
        state["include_references"] = True
    
    return state


def build_chat_graph():
    """构建增强的聊天图"""
    graph = StateGraph(ChatState)
    
    # 添加所有节点
    graph.add_node("initialize", initialize_state)
    graph.add_node("intent_classification", classify_intent)
    graph.add_node("code_analysis", analyze_cairo_code)
    graph.add_node("error_analysis", analyze_error)
    graph.add_node("knowledge_retrieval", retrieve_knowledge)
    graph.add_node("research_planning", plan_research)
    graph.add_node("research_execution", execute_research_stage)
    graph.add_node("research_synthesis", synthesize_research)
    graph.add_node("research_conclusion", conclude_research)
    graph.add_node("answer_generation", generate_answer)
    
    # 设置入口点
    graph.set_entry_point("initialize")
    
    # 设置基本流程
    graph.add_edge("initialize", "intent_classification")
    
    # 意图分类后的条件路由
    graph.add_conditional_edges(
        "intent_classification",
        route_after_intent_classification,
        {
            "error_analysis": "error_analysis",
            "code_analysis": "code_analysis", 
            "research_planning": "research_planning",
            "knowledge_retrieval": "knowledge_retrieval"
        }
    )
    
    # 代码分析后的条件路由
    graph.add_conditional_edges(
        "code_analysis",
        route_after_code_analysis,
        {
            "error_analysis": "error_analysis",
            "research_planning": "research_planning",
            "knowledge_retrieval": "knowledge_retrieval"
        }
    )
    
    # 错误分析后的条件路由
    graph.add_conditional_edges(
        "error_analysis",
        route_after_error_analysis,
        {
            "research_planning": "research_planning",
            "answer_generation": "answer_generation"
        }
    )
    
    # 知识检索后的条件路由
    graph.add_conditional_edges(
        "knowledge_retrieval",
        route_after_knowledge_retrieval,
        {
            "research_planning": "research_planning",
            "answer_generation": "answer_generation"
        }
    )
    
    # 研究规划后进入研究执行
    graph.add_edge("research_planning", "research_execution")
    
    # 研究执行的条件路由（支持循环）
    graph.add_conditional_edges(
        "research_execution",
        route_research_execution,
        {
            "research_execution": "research_execution",  # 循环
            "knowledge_retrieval": "knowledge_retrieval",
            "research_synthesis": "research_synthesis",
            "research_conclusion": "research_conclusion",
            "answer_generation": "answer_generation"
        }
    )
    
    # 研究综合后的路由
    graph.add_conditional_edges(
        "research_synthesis",
        should_continue_research,
        {
            "research_execution": "research_execution",
            "answer_generation": "answer_generation"
        }
    )
    
    # 研究结论后生成答案
    graph.add_edge("research_conclusion", "answer_generation")
    
    # 答案生成后结束
    graph.add_edge("answer_generation", END)
    
    # 编译图并设置递归限制
    return graph.compile(
        checkpointer=None,
        interrupt_before=None,
        interrupt_after=None,
        debug=False
    )
